"""
API模型定义，使用Pydantic定义请求和响应模型
"""
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field

class TextGenerationRequest(BaseModel):
    """文本生成请求模型"""
    prompt: str = Field(..., description="输入提示文本")
    max_new_tokens: Optional[int] = Field(None, description="最大生成的新token数量")
    temperature: Optional[float] = Field(None, description="温度参数，控制随机性，范围0-2.0")
    top_p: Optional[float] = Field(None, description="核采样参数，范围0-1.0")
    top_k: Optional[int] = Field(None, description="top-k采样参数")
    
    class Config:
        schema_extra = {
            "example": {
                "prompt": "请介绍一下ROCm是什么",
                "max_new_tokens": 256,
                "temperature": 0.7,
                "top_p": 0.9,
                "top_k": 50
            }
        }

class TextGenerationResponse(BaseModel):
    """文本生成响应模型"""
    prompt: str = Field(..., description="原始输入提示文本")
    generated_text: str = Field(..., description="完整生成的文本（包含原始提示）")
    response: str = Field(..., description="模型的回复（不包含原始提示）")
    metadata: Dict[str, Any] = Field(..., description="生成的元数据信息")

class BatchGenerationRequest(BaseModel):
    """批量文本生成请求模型"""
    prompts: List[str] = Field(..., description="输入提示文本列表")
    max_new_tokens: Optional[int] = Field(None, description="最大生成的新token数量")
    temperature: Optional[float] = Field(None, description="温度参数，控制随机性")
    top_p: Optional[float] = Field(None, description="核采样参数")
    top_k: Optional[int] = Field(None, description="top-k采样参数")
    
    class Config:
        schema_extra = {
            "example": {
                "prompts": ["请介绍一下ROCm是什么", "ROCm和CUDA有什么区别"],
                "max_new_tokens": 256,
                "temperature": 0.7,
                "top_p": 0.9,
                "top_k": 50
            }
        }

class BatchGenerationResponse(BaseModel):
    """批量文本生成响应模型"""
    results: List[TextGenerationResponse] = Field(..., description="生成结果列表")

class ModelInfoResponse(BaseModel):
    """模型信息响应模型"""
    model_name: str = Field(..., description="模型名称")
    model_path: str = Field(..., description="模型路径")
    device: str = Field(..., description="运行设备")
    precision: str = Field(..., description="精度")
    max_length: int = Field(..., description="最大输入长度")
    max_new_tokens: int = Field(..., description="默认最大生成token数")